#!/bin/env python
from netCDF4 import Dataset
import numpy as np
import math

# ENGPI computation using a renalayis dataset

#--Directory Setting
input_reanalysis_dir = "./input_data/reanalysis" # input directory for reanalysis data
input_besttrack_dir = "./input_data/besttrack" # input directory for tropical cyclone tracks
output_dir = "./output_data" # output directory
output_file1 = "%s/engpi.nc" % (output_dir) # output: monthly ENGPI data
output_file2 = "%s/engpi_climmonthly.nc" % (output_dir) # output: climatological monthly ENGPI data

#--Set Parameters
##--DGPI coefficients
coef={
  'mpi'     :3.0,
  'wshear'  :-2.0,
  'avor'    :2.0,
  'rh'      :3.0,
}
undef = -9.99E33 # undefined value

#--Read inputs
##--potential intensity
f1 = Dataset("%s/mpi.nc" % (input_reanalysis_dir), "r")
mpi = f1.variables['mpi']

##- relative humidity at 600 hPa
f2 = Dataset("%s/rh600.nc" % (input_reanalysis_dir), "r")
rh = f2.variables['rh600']

##--absolute vorticity at 850 hPa
f3 = Dataset("%s/avor850a.nc" % (input_reanalysis_dir), "r")
avor = f3.variables['avor850a']

##--vertical wind shear
f4 = Dataset("%s/wshear.nc" % (input_reanalysis_dir), "r")
wshear = f4.variables['wshear']

##--relative SST anomaly
f5 = Dataset("%s/sstanm.nc" % (input_reanalysis_dir), "r")
ssta = f5.variables['sstanm']

tmax,jmax,imax = np.shape(mpi[:])
print ("dimension=",tmax,jmax,imax)

lon_in = f5.variables['lon']
lat_in = f5.variables['lat']
times_in = f5.variables['time']

lons,lats = np.meshgrid(lon_in[:],lat_in[:])
lats_tile = np.tile(lats,(tmax,1,1))

#-----Compute DGPI---------------------------
mpi_term = (mpi[:]/70.0)**(coef['mpi'])
wshear_term = (1.0 + 0.1*wshear[:])**(coef['wshear'])
rh_term = (rh[:]/50)**(coef['rh'])
avor_term = (np.absolute(100000*avor[:]))**(coef['avor'])
engpi = mpi_term * wshear_term * rh_term * avor_term 

##-----Post Process---------------------------
##--Fill zero value where latitude is smaller than 2 degree
neg_lat_index = np.where((lats_tile<2)&(lats_tile>-2))
engpi[neg_lat_index] = 0.0
mpi_term[neg_lat_index] = 0.0
wshear_term[neg_lat_index] = 0.0
rh_term[neg_lat_index] = 0.0
avor_term[neg_lat_index] = 0.0
 
##--Mask over land surface
engpi = np.ma.masked_where(ssta[:].mask==True, engpi)
mpi_term = np.ma.masked_where(ssta[:].mask==True, mpi_term)
wshear_term = np.ma.masked_where(ssta[:].mask==True, wshear_term)
rh_term = np.ma.masked_where(ssta[:].mask==True, rh_term)
avor_term = np.ma.masked_where(ssta[:].mask==True, avor_term)

#print engpi,np.shape(engpi),type(engpi)

#-----Output---------------------------------
#-----Monthly
f6 = Dataset(output_file1, "w", format="NETCDF3_CLASSIC")

#dimmensions
time = f6.createDimension('time', None)
lat = f6.createDimension('lat',jmax)
lon = f6.createDimension('lon',imax)

#variables
times = f6.createVariable('time','f8',('time',))
latitudes = f6.createVariable('lat','f4', ('lat',))
longitudes = f6.createVariable('lon','f4', ('lon',))

engpi_out = f6.createVariable('engpi','f4',('time','lat','lon'),fill_value=undef)
mpi_out = f6.createVariable('mpi','f4',('time','lat','lon'),fill_value=undef)
wshear_out = f6.createVariable('wshear','f4',('time','lat','lon'),fill_value=undef)
rh_out = f6.createVariable('rh','f4',('time','lat','lon'),fill_value=undef)
avor_out = f6.createVariable('avor','f4',('time','lat','lon'),fill_value=undef)
    
#attributtes
f6.description = 'DGPI'
latitudes.units = lat_in.units
latitudes.axis= 'Y'
longitudes.units = lon_in.units
longitudes.axis= 'X'
times.units = times_in.units
times.calendar = times_in.calendar
times.cargtesian_axis = "T"

engpi_out.long_name = "Monthly DGPI"
mpi_out.long_name = "MPI Term"
wshear_out.long_name = "Vertical Wind Shear Term"
rh_out.long_name = "Relative Humidity Term"
avor_out.long_name = "Absolute Vorticity Term"

latitudes[:] = lat_in[:]
longitudes[:] = lon_in[:]
times[:] = times_in[:]

engpi_out[:] = engpi[:]
mpi_out[:] = mpi_term[:]
wshear_out[:] = wshear_term[:]
rh_out[:] = rh_term[:]
avor_out[:] = avor_term[:]

#-----Climatological Monthly
f7 = Dataset(output_file2, "w", format="NETCDF3_CLASSIC")

#dimmensions
time = f7.createDimension('time', 12)
lat = f7.createDimension('lat',jmax)
lon = f7.createDimension('lon',imax)

#variables
times = f7.createVariable('time','f8',('time',))
latitudes = f7.createVariable('lat','f4', ('lat',))
longitudes = f7.createVariable('lon','f4', ('lon',))

engpi_out = f7.createVariable('engpi','f4',('time','lat','lon'),fill_value=undef)
mpi_out = f7.createVariable('mpi','f4',('time','lat','lon'),fill_value=undef)
wshear_out = f7.createVariable('wshear','f4',('time','lat','lon'),fill_value=undef)
rh_out = f7.createVariable('rh','f4',('time','lat','lon'),fill_value=undef)
avor_out = f7.createVariable('avor','f4',('time','lat','lon'),fill_value=undef)
    
#attributtes
f7.description = 'DGPI'
latitudes.units = lat_in.units
latitudes.axis= 'Y'
longitudes.units = lon_in.units
longitudes.axis= 'X'
times.units = times_in.units
times.calendar = times_in.calendar
times.cargtesian_axis = "T"

engpi_out.long_name = "Climatological Monthly ENGPI"
mpi_out.long_name = "MPI Term"
wshear_out.long_name = "Vertical Wind Shear Term"
rh_out.long_name = "Relative Humidity Term"
avor_out.long_name = "Absolute Vorticity Term"

latitudes[:] = lat_in[:]
longitudes[:] = lon_in[:]
times[:] = times_in[0:12]

engpi_out[:] = engpi[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
mpi_out[:] = mpi_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
wshear_out[:] = wshear_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
rh_out[:] = rh_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
avor_out[:] = avor_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)

#-----Finalize
f1.close()
f2.close()
f3.close()
f4.close()
f5.close()
f6.close()
f7.close()
